#!/usr/bin/env python
#-*- coding:utf-8 -*-
# Author: LiuHuan
# Datetime: 2020/11/13 11:04

import json
import pandas as pd
import sys
import os
os.chdir(sys.path[0])

if __name__ == '__main__':

    ## 设置生成训练数据的路径
    path = r"E:\LiuHuan\Projects\bert\work\2_layer_cla_cscd\S\data_en"

    ## 先把 id2label.json 和 label2id.json都拷进上面的路径
    with open(os.path.join(path,'id2label.json'),'r',encoding='utf-8') as f:
        id2label = json.load(f)


    data_train = pd.DataFrame(columns=('id', 'title', 'abstract', 'en_abstract'))
    data_test = pd.DataFrame(columns=('id', 'title', 'abstract', 'en_abstract'))

    ##设置读取数据的路径
    path_read = r"E:\LiuHuan\Projects\bert\work\2_layer_cla_cscd\S\data_csv"

    print(id2label)
    for id,label in id2label.items():
        df = pd.read_csv(os.path.join(path_read,label+'.csv'))
        if len(df) > 5000:
            df = df[:5000]

        df_train = df[:int(len(df) * 0.9)]
        df_test = df[int(len(df) * 0.9):]

        abstracts = []
        for j in range(len(df_train)):
            if not pd.isna(df_train.iloc[j]['en_abstract']):
                abst = df_train.iloc[j]['en_abstract'].strip().replace('\r', '').replace('\n', '')
                if abst != '':
                    abstracts.append(abst)
        with open(os.path.join(path, 'train_temp.tsv'), 'a', encoding='utf-8') as f:
            for abst in abstracts:
                f.write(id + '\t' + abst + '\n')
        print(len(abstracts))
        abstracts = []
        for j in range(len(df_test)):
            if not pd.isna(df_test.iloc[j]['en_abstract']):
                abst = df_test.iloc[j]['en_abstract'].strip().replace('\r', '').replace('\n', '')
                if abst != '':
                    abstracts.append(abst)

        with open(os.path.join(path, 'test.tsv'), 'a', encoding='utf-8') as f:
            for abst in abstracts:
                f.write(id + '\t' + abst + '\n')

        with open(os.path.join(path, 'dev.tsv'), 'a', encoding='utf-8') as f:
            for abst in abstracts:
                f.write(id + '\t' + abst + '\n')

        print(len(df))
        print(label, ' Done')

    df_train = pd.read_csv(os.path.join(path, 'train_temp.tsv'), sep='\t', names=['label', 'Sentence'])

    print(len(df_train))

    df_train = df_train.sample(frac=1).reset_index(drop=True)
    df_train.to_csv(os.path.join(path, 'train.tsv'), sep='\t', header=False, index=False)